import sys
from PyQt5.QtWidgets import QApplication, QVBoxLayout, QWidget, QPushButton
from PyQt5.QtCore import Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.animation as animation

class BubbleSortAnimation(QWidget):
    def __init__(self):
        super().__init__()

        self.initUI()

        # 初始化记录排序过程中各个步骤的列表
        self.steps = []

    def initUI(self):
        """
        初始化用户界面，包括窗口、布局、按钮和绘图区域。
        """
        # 设置窗口标题和大小
        self.setWindowTitle('Bubble Sort Animation')
        self.setGeometry(100, 100, 800, 600)

        # 创建垂直布局
        self.layout = QVBoxLayout()

        # 创建matplotlib的Figure对象和FigureCanvas对象
        self.figure = Figure()
        self.canvas = FigureCanvas(self.figure)
        self.layout.addWidget(self.canvas)

        # 创建并添加启动按钮
        self.start_button = QPushButton('Start Sort', self)
        self.start_button.clicked.connect(self.start_sort)
        self.layout.addWidget(self.start_button)

        # 将布局应用到窗口
        self.setLayout(self.layout)

        # 初始绘制柱形图，显示初始状态
        self.plot_bars([5, 3, 8, 4, 2], -1, -1)

    def start_sort(self):
        """
        开始冒泡排序的动画，并记录排序过程中的每一步状态。
        """
        # 重置排序数组
        self.numbers = [5, 3, 8, 4, 2]

        # 记录排序过程
        self.record_sorting_steps()

        # 启动动画，每500毫秒更新一次图形
        self.ani = animation.FuncAnimation(self.figure, self.update_plot, frames=len(self.steps), interval=500, repeat=False, blit=False)

        # 添加动画结束的事件处理，绘制最终状态
        self.ani.event_source.add_callback(self.plot_final_state)

        # 绘制第一帧的图形
        self.canvas.draw()

    def record_sorting_steps(self):
        """
        记录冒泡排序的每一步状态到self.steps列表中。
        """
        # 清空之前记录的步骤
        self.steps = []

        n = len(self.numbers)
        for i in range(n):
            for j in range(n - i - 1):
                # 记录当前状态
                self.steps.append((self.numbers[:], j, j + 1))
                if self.numbers[j] > self.numbers[j + 1]:
                    # 交换元素
                    self.numbers[j], self.numbers[j + 1] = self.numbers[j + 1], self.numbers[j]
                    # 记录交换后的状态
                    self.steps.append((self.numbers[:], j, j + 1))

    def update_plot(self, frame):
        """
        根据动画的帧数更新柱形图。
        """
        # 获取当前步骤的状态
        numbers, j, next_j = self.steps[frame]
        # 绘制当前步骤的柱形图
        self.plot_bars(numbers, j, next_j)

    def plot_bars(self, numbers, j, next_j):
        """
        绘制柱形图，用红色标记当前比较的两个柱子。
        """
        # 清除之前的绘图
        self.figure.clear()
        ax = self.figure.add_subplot(111)

        # 设置柱形的颜色，默认是蓝色，比较的两个元素为红色
        bar_colors = ['b'] * len(numbers)
        if j != -1 and next_j != -1:
            bar_colors[j] = 'r'
            bar_colors[next_j] = 'r'

        # 绘制柱形图
        ax.bar(range(len(numbers)), numbers, color=bar_colors)
        self.canvas.draw()

    def plot_final_state(self):
        """
        绘制排序完成后的最终状态，确保最后两个柱子正确显示为蓝色。
        """
        # 获取最终排序完成后的状态
        final_numbers = self.steps[-1][0]
        # 绘制最终状态的柱形图
        self.plot_bars(final_numbers, -1, -1)

if __name__ == '__main__':
    # 创建Qt应用
    app = QApplication(sys.argv)

    # 创建并显示BubbleSortAnimation窗口
    ex = BubbleSortAnimation()
    ex.show()

    # 进入Qt应用的主循环
    sys.exit(app.exec_())
